# -*- coding: utf-8 -*-
# @Time    : 2022/2/5 9:36 上午
# @Author  : DIRICHLET
# @Email   : 511827625@qq.com
# @File    : ConditionalAutoEncoder.py
# @Software: PyCharm
# @ Better Late Than Never!

import numpy as np
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras import layers

class ConditionalAutoEncoder(Model):
    def __init__(self,
                 inputSize,
                 encodeHiddenSize,
                 zSize,
                 decodeHiddenSize,
                 ySize
                 ):
        super().__init__()
        self.encoder=Encoder(encodeHiddenSize,zSize)
        self.decoder=Decoder(decodeHiddenSize,zSize,ySize,inputSize)
        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = tf.keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.kl_loss_tracker = tf.keras.metrics.Mean(name="kl_loss")


    def call(self,x,y=None,mode='train'):
        if mode=='train':
            #如果是训练模式 就先从x获得z
            _,_,z=self.encoder(x)
        else:
            # 不是训练模式 就随机生成z 标准正态分布
            # np.random.randn(show_num, 2)
            # 不需要x 但是需要onehot y
            z=tf.random.normal(tf.stack([tf.shape(x)[0], self.encoder.zSize]), 0, 1, dtype=tf.float32)

        if y is None:
            raise Exception('erro y is None')

        nz=tf.concat([z,y],1)
        #=self.decode(nz)
        return self.decoder(nz)

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    def train_step(self, data):
        x,y=data
        data=tf.concat([x,y],1)
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            nz = tf.concat([z, y], 1)
            reconstruction = self.decoder(nz)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    tf.keras.losses.binary_crossentropy(x, reconstruction)
                )
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
            total_loss = reconstruction_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }


class Encoder(Model):
    def __init__(self,hiddenSize,zSize):
        super(Encoder, self).__init__()
        self.hiddenSize=hiddenSize
        self.hidden=layers.Dense(self.hiddenSize,activation='relu')
        self.zSize=zSize
        self.zMean=layers.Dense(self.zSize,activation='relu')
        self.zSigma=layers.Dense(self.zSize,activation='relu')
        # 此处为log(sigma^2) 映射到-inf，+inf



    def  call(self, x):
        x=self.hidden(x)
        zMean=self.zMean(x)
        zSigma=self.zSigma(x)
        eps = tf.random.normal(tf.stack([tf.shape(x)[0], self.zSize]), 0, 1, dtype=tf.float32)
        # 标准正态分布
        z = tf.add(zMean, tf.multiply(tf.exp(0.5*zSigma), eps))

        return zMean,zSigma,z


class Decoder(Model):
    def __init__(self,
                 hiddenSize,
                 zSize,
                 ySize,
                 outSize):
        super(Decoder, self).__init__()
        self.hiddenSize = hiddenSize
        self.hidden = layers.Dense(self.hiddenSize, activation='relu')
        self.conditionalSize = zSize+ySize
        self.outSize=outSize
        self.out = layers.Dense(self.outSize, activation='relu')

    def call(self, x):
        x = self.hidden(x)
        x = self.out(x)
        return x





